#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06

"""
【节点名称】：
    COCODatasetLoaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    std_msgs::Boolean （是否输出下一张图像）
    spirecv_msgs::EvaluationResult （算法评价结果）
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
    std_msgs::Boolean （评价任务结束时，输出关闭所有节点指令）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown_emit (bool): 评价任务结束时，是否输出关闭所有节点指令
【节点参数】：
    coco_root_dir (str): 数据集根目录
    coco_gt_file (str): 数据集COCO类型真值
    img_total (int): 数据集图像数量
【备注】：
    无
"""

import threading
import time
import os
import platform
import json
from collections import defaultdict
import cv2
import argparse
from typing import Union
from queue import Queue
from spirems import Subscriber, Publisher, cvimg2sms, def_msg, QoS, BaseNode, get_extra_args
import uuid


class COCODatasetLoaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.sms_shutdown_emit = self.get_param("sms_shutdown_emit", True)
        self.coco_root_dir = self.get_param("coco_root_dir", r"D:\dataset\val2017")
        self.coco_gt_file = self.get_param("coco_gt_file", r"D:\dataset\annotations\instances_val2017.json")
        self.img_total = self.get_param("img_total", 5000)
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True
    
        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.queue_pool.append(self.img_i_queue)
        self.img_i = 0

        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )

        evaluate_url = '/SpireView/{}/TaskResults'.format(self.client_id)
        self._evaluate_reader = Subscriber(
            evaluate_url, 'spirecv_msgs::EvaluationResult', self.evaluate_callback,
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        self._evaluate_writer = Publisher(
            '/SpireView/EvaluationResult', 'spirecv_msgs::EvaluationResult',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._shutdown_writer = Publisher(
            '/' + job_name + '/shutdown', 'std_msgs::Boolean',
            ip=ip, port=port, qos=QoS.Reliability
        )

        self._load_dataset()
        self.start()

    def release(self):
        BaseNode.release(self)
        self._next_reader.kill()
        self._evaluate_reader.kill()
        self._evaluate_writer.kill()
        self._image_writer.kill()
        self._shutdown_writer.kill()
        self._next_reader.join()
        self._evaluate_reader.join()
        self._evaluate_writer.join()
        self._image_writer.join()
        self._shutdown_writer.join()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.img_keys):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1

    def evaluate_callback(self, msg):
        print(msg)
        self._evaluate_writer.publish(msg)
        if self.sms_shutdown_emit:
            msg = def_msg('std_msgs::Boolean')
            msg['data'] = True
            self._shutdown_writer.publish(msg)

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img_info = self.imgs[self.img_keys[img_i]]
            # img_id = img_info['id']
            file_name = os.path.basename(img_info['file_name'])
            img = cv2.imread(str(os.path.join(self.coco_root_dir, file_name)))

            assert img.shape[0] == img_info['height'] and img.shape[1] == img_info['width']
            # print(img.shape)
            if self.b_use_shm:
                msg = self._image_writer.cvimg2sms_mem(img)
            else:
                msg = cvimg2sms(img, 'png')

            msg['img_id'] = img_i
            msg['img_total'] = self.img_total
            msg['file_name'] = file_name
            msg['client_id'] = self.client_id
            self._image_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))

    def _load_dataset(self):
        annotation_file = self.coco_gt_file
        tic = time.time()
        with open(annotation_file, 'r') as f:
            dataset = json.load(f)
        assert isinstance(dataset, dict), 'annotation file format {} not supported'.format(type(dataset))
        print('Done (t={:0.2f}s)'.format(time.time() - tic))
        self.dataset = dataset
        self._create_index()
        self.img_keys = list(self.imgs.keys())

    def _create_index(self):
        # create index
        print('creating index...')
        anns, cats, imgs = {}, {}, {}
        imgToAnns, catToImgs = defaultdict(list), defaultdict(list)
        if 'annotations' in self.dataset:
            for ann in self.dataset['annotations']:
                imgToAnns[ann['image_id']].append(ann)
                anns[ann['id']] = ann

        if 'images' in self.dataset:
            for img in self.dataset['images']:
                imgs[img['id']] = img

        if 'categories' in self.dataset:
            for cat in self.dataset['categories']:
                cats[cat['id']] = cat

        if 'annotations' in self.dataset and 'categories' in self.dataset:
            for ann in self.dataset['annotations']:
                catToImgs[ann['category_id']].append(ann['image_id'])

        print('index created!')
        # create class members
        self.anns = anns
        self.imgToAnns = imgToAnns
        self.catToImgs = catToImgs
        self.imgs = imgs
        self.cats = cats


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    loader = COCODatasetLoaderNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    loader.launch_next()
    loader.join()
